#!/bin/bash

# 设置模型路径
export ModelPath="$1"

# 设置使用卡的数量
export WorldSize="$2"

# 设置模型名称
export ModelName="$3"

# 检查是否提供了足够的参数
if [ -z "$ModelPath" ] || [ -z "$WorldSize" ] || [ -z "$ModelName" ]; then
    echo "使用方法: $0 <模型路径> <卡数> <模型名称>"
    exit 1
fi

# 验证卡数是否为数字
if ! [[ "$WorldSize" =~ ^[0-9]+$ ]]; then
    echo "错误: 卡数必须是一个有效的数字!"
    exit 1
fi

# 生成设备ID列表
generate_device_ids() {
    local world_size=$1
    local ids=""
    for ((i=0; i<world_size; i++)); do
        if [ $i -eq 0 ]; then
            ids="$i"
        else
            ids="$ids,$i"
        fi
    done
    echo "$ids"
}

device_ids=$(generate_device_ids $WorldSize)


# 获取当前服务器IP地址
ip_address=$(hostname -I | awk '{print $1}')

# 输出参数
echo "- worldSize: $WorldSize"
echo "- npuDeviceIds: [[$device_ids]]"
echo "- ip_address: $ip_address"
echo "- ModelName: $ModelName"

# 启动服务
# 检测硬件设备类型
if lspci | grep -q "d500"; then
    echo "检测到 d500 设备，启动 310 版本"

    # 设置docker compose的启动参数
    export ServicePath=$(pwd)/310/service-run.sh
    export ConfigPath=$(pwd)/310/config.json

    # 设置大模型名称
    sed -i "s/\"modelName\"[[:space:]]*:[[:space:]]*\"[^\"]*\"/\"modelName\" : \"$ModelName\"/g" $(pwd)/310/config.json

    # 自动设置IP
    sed -i "s/\"ipAddress\"[[:space:]]*:[[:space:]]*\"[^\"]*\"/\"ipAddress\" : \"$ip_address\"/g" $(pwd)/310/config.json

    # 设置卡的数量/ID
    sed -i "s|\"worldSize\" *: *[0-9]*,|\"worldSize\": $WorldSize,|" $(pwd)/310/config.json
    sed -i "s/\"npuDeviceIds\"[[:space:]]*:[[:space:]]*\[\[[0-9,]*\]\],/\"npuDeviceIds\" : [[$device_ids]],/" $(pwd)/310/config.json

    # 执行脚本添加可执行权限
    chmod +x $(pwd)/310/service-run.sh

    # 启动服务
    cd "$(pwd)/310" && docker-compose up -d

# 检测硬件设备类型
elif lspci | grep -q "d802"; then
    echo "检测到 d802 设备，启动 910 版本"

    # 设置docker compose的启动参数
    export ServicePath=$(pwd)/910/service-run.sh  # 统一变量名为ServicePath
    export ConfigPath=$(pwd)/910/config.json

    # 设置大模型名称
    sed -i "s/\"modelName\"[[:space:]]*:[[:space:]]*\"[^\"]*\"/\"modelName\" : \"$ModelName\"/g" $(pwd)/910/config.json

    # 自动设置IP
    sed -i "s/\"ipAddress\"[[:space:]]*:[[:space:]]*\"[^\"]*\"/\"ipAddress\" : \"$ip_address\"/g" $(pwd)/910/config.json

    # 设置卡的数量
    sed -i "s|\"worldSize\" *: *[0-9]*,|\"worldSize\": $WorldSize,|" $(pwd)/910/config.json
    sed -i "s/\"npuDeviceIds\"[[:space:]]*:[[:space:]]*\[\[[0-9,]*\]\],/\"npuDeviceIds\" : [[$device_ids]],/" $(pwd)/910/config.json

    # 执行脚本添加可执行权限
    chmod +x $(pwd)/910/service-run.sh

    # 启动服务
    cd "$(pwd)/910" && docker-compose up -d

else
    echo "错误：未检测到 d500 或 d802 设备"
    exit 1
fi